import torch
from torchvision import transforms
from PIL import Image

import logging

from . import get_model

def extract(img_path: str):

    # load the best model with PCA (trained by our SFRS)
    model = get_model.get()

    # read image
    img = Image.open(img_path).convert(
        "RGB"
    )  # modify the image path according to your need
    transformer = transforms.Compose(
        [
            transforms.Resize((480, 640)),  # (height, width)
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.48501960784313836, 0.4579568627450961, 0.4076039215686255],
                std=[0.00392156862745098, 0.00392156862745098, 0.00392156862745098],
            ),
        ]
    )
    img = transformer(img)

    # use GPU (optional)
    model = model.cuda()
    img = img.cuda()

    # extract descriptor (4096-dim)
    with torch.no_grad():
        des = model(img.unsqueeze(0))[0]
    des = des.cpu().numpy()

    return des
